package com.ppch.wuwamanus.advisors;

import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.ppch.wuwamanus.model.entity.ChatHistory;
import com.ppch.wuwamanus.model.enums.ChatHistoryMessageTypeEnum;
import com.ppch.wuwamanus.service.ChatHistoryService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * Package:WuWaManus
 * ClassName:DatabaseChatMemory
 *
 * @Author 泡泡茶壶
 * @Create 2025/9/19 15:35
 * @Version 1.0
 * Description:
 *  基于MySql的对话记忆持久化
 */
@Slf4j
@Component
public class DatabaseChatMemory implements ChatMemory {

    /**
     * 对话记录 ChatHistoryService
     */
    private final ChatHistoryService chatHistoryService;

    /**
     * 当前用户ID
     */
    private Long currentUserId;

    /**
     * 构造函数，初始化 ChatHistoryService。
     * @param chatHistoryService ChatHistoryService
     */
    public DatabaseChatMemory(ChatHistoryService chatHistoryService) {
        this.chatHistoryService = chatHistoryService;
    }

    /**
     * 设置当前用户ID
     *
     * @param userId 用户ID
     */
    public void setCurrentUserId(Long userId) {
        this.currentUserId = userId;
    }

    /**
     * 添加对话记录
     * 1. 将消息列表存储到数据库中
     *
     * @param conversationId 会话ID
     * @param messages       消息列表
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        Long chatId = Long.valueOf(conversationId);
        for (Message message : messages) {
            MessageType messageType = message.getMessageType();
            if(messageType == MessageType.USER){
                String userText = message.getText();
                if(StrUtil.isNotBlank(userText)){
                    chatHistoryService.addChatHistory(chatId,userText, ChatHistoryMessageTypeEnum.USER.getValue(), currentUserId);
                }
                log.warn("DatabaseChatMemory 添加用户消息到数据库中，当前会话ID: {}, 当前用户ID: {}, 用户消息: {}", chatId, currentUserId, userText);
            }else if(messageType == MessageType.ASSISTANT){
                String text = message.getText();
                if(StrUtil.isNotBlank(text)){
                    chatHistoryService.addChatHistory(chatId,text, ChatHistoryMessageTypeEnum.AI.getValue(), currentUserId);
                }
                log.warn("ASSISTANT消息: {}", text);
            }else if(messageType == MessageType.SYSTEM){
                String text = message.getText();
                if(StrUtil.isNotBlank(text)){
                    chatHistoryService.addChatHistory(chatId,text, ChatHistoryMessageTypeEnum.AI.getValue(), currentUserId);
                }
                log.warn("SYSTEM消息: {}", text);
            }else if(messageType == MessageType.TOOL){
                ToolResponseMessage toolResponseMessage = (ToolResponseMessage) message;
                String toolResponseText = toolResponseMessage.getResponses().stream()
                        .map(toolResponse -> "工具执行结果如下: " + toolResponse.responseData()).collect(Collectors.joining("\n"));
                chatHistoryService.addChatHistory(chatId,toolResponseText, ChatHistoryMessageTypeEnum.AI.getValue(), currentUserId);
                log.warn("TOOL消息: {}", toolResponseText);
            }
        }
    }

    /**
     * 根据会话ID获取对应的对话历史记录
     * @param conversationId 会话ID
     * @param lastN 获取最近几条对话记录
     * @return 对话历史记录列表 List<Message>
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        Long chatId = Long.valueOf(conversationId);
        // 根据会话ID查询对话历史记录
        LambdaQueryWrapper<ChatHistory> queryWrapper = new LambdaQueryWrapper<>();
        queryWrapper.eq(ChatHistory::getChatId, chatId);
        List<ChatHistory> chatHistoryList = chatHistoryService.list(queryWrapper);
        List<Message> messages = new ArrayList<>();
        // 将对话历史记录转换为消息列表（需要根据消息类型进行转换）
        // ai消息：转换为 AssistantMessage; user消息：转换为 UserMessage
        for (ChatHistory chatHistory : chatHistoryList) {
            if (chatHistory.getMessageType() == ChatHistoryMessageTypeEnum.AI.getValue()) {
                messages.add(new AssistantMessage(chatHistory.getMessage()));
            } else {
                messages.add(new UserMessage(chatHistory.getMessage()));
            }
        }
        // 返回最近N条对话记录
        return messages != null ? messages.stream().skip(Math.max(0, messages.size() - lastN)).toList() : List.of();
    }

    /**
     * 清空指定会话ID的对话历史记录
     * @param conversationId
     */
    @Override
    public void clear(String conversationId) {
        Long chatId = Long.valueOf(conversationId);
        boolean b = chatHistoryService.removeById(chatId);
        if(b){
            log.info("DatabaseChatMemory 清空指定会话ID的对话历史记录成功，当前会话ID: {}", chatId);
        }else {
            log.error("DatabaseChatMemory 清空指定会话ID的对话历史记录失败，当前会话ID: {}", chatId);
        }
    }
}
